#include "Shader.h"
#include "Util.h"

namespace TestDX
{

Shader::Shader() :
	_pVertexShader(NULL),
	_pPixelShader(NULL),
	_pLayout(NULL),
	_pSamplerState(NULL),
	_pMatrixBuffer(NULL),
	_pCameraBuffer(NULL),
	_pLightBuffer(NULL)
{
}


Shader::~Shader()
{
	Shutdown();
}

bool Shader::Init(ID3D11Device * device, HWND hwnd)
{
	return InitShader(device, hwnd, L"light.vs", L"light.ps");
}

bool Shader::InitShader(ID3D11Device * device,
						HWND hwnd,
						WCHAR * vsName,
						WCHAR * psName)
{
	ID3D10Blob * errorMsg = NULL;

	// Compile Vertex Shader
	ID3D10Blob * vsBuffer = NULL;
	HRESULT hr = D3DX11CompileFromFile(vsName,
									   NULL,
									   NULL,
									   "LightVertexShader",
									   "vs_5_0",
									   D3D10_SHADER_ENABLE_STRICTNESS,
									   0,
									   NULL,
									   &vsBuffer,
									   &errorMsg,
									   NULL);

	if(FAILED(hr))
	{
		if(errorMsg)
			OutputShaderErrorMessage(errorMsg, hwnd, vsName);
		else
			// Empty errorMsg means it could not find the file
			MessageBox(hwnd, vsName, L"Missing Shader File", MB_OK);
		
		return false;
	}

	// Compile Pixel Shader
	ID3DBlob * psBuffer = NULL;
	hr = D3DX11CompileFromFile(psName,
							   NULL,
							   NULL,
							   "LightPixelShader",
							   "ps_5_0",
							   D3D10_SHADER_ENABLE_STRICTNESS,
							   0,
							   NULL,
							   &psBuffer,
							   &errorMsg,
							   NULL);
	
	if(FAILED(hr))
	{
		if(errorMsg)
			OutputShaderErrorMessage(errorMsg, hwnd, psName);
		else
			MessageBox(hwnd, psName, L"Missing Shader File", MB_OK);
		
		return false;
	}

	// Create Vertex Shader from buffer
	HR(device->CreateVertexShader(vsBuffer->GetBufferPointer(),
								  vsBuffer->GetBufferSize(),
								  NULL,
								  &_pVertexShader));

	HR(device->CreatePixelShader(psBuffer->GetBufferPointer(),
								 psBuffer->GetBufferSize(),
								 NULL,
								 &_pPixelShader));

	// Create Vertex input layout description
	// It has to match VertexType struct in Model class and in the shader
	D3D11_INPUT_ELEMENT_DESC ied[3];
	ZeroMemory(&ied, sizeof(D3D11_INPUT_ELEMENT_DESC) * 3);

	ied[0].SemanticName = "POSITION";
	ied[0].SemanticIndex = 0;
	ied[0].Format = DXGI_FORMAT_R32G32B32_FLOAT;
	ied[0].InputSlot = 0;
	ied[0].AlignedByteOffset = 0;
	ied[0].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
	ied[0].InstanceDataStepRate = 0;

	ied[1].SemanticName = "TEXCOORD";
	ied[1].SemanticIndex = 0;
	ied[1].Format = DXGI_FORMAT_R32G32_FLOAT;
	ied[1].InputSlot = 0;
	ied[1].AlignedByteOffset = D3D11_APPEND_ALIGNED_ELEMENT;
	ied[1].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
	ied[1].InstanceDataStepRate = 0;

	ied[2].SemanticName = "NORMAL";
	ied[2].SemanticIndex = 0;
	ied[2].Format = DXGI_FORMAT_R32G32B32_FLOAT;
	ied[2].InputSlot = 0;
	ied[2].AlignedByteOffset = D3D11_APPEND_ALIGNED_ELEMENT;
	ied[2].InputSlotClass = D3D11_INPUT_PER_VERTEX_DATA;
	ied[2].InstanceDataStepRate = 0;

	unsigned int numElements = sizeof(ied) / sizeof(ied[0]);

	HR(device->CreateInputLayout(ied,
								 numElements,
								 vsBuffer->GetBufferPointer(),
								 vsBuffer->GetBufferSize(),
								 &_pLayout));

	// Release vs and ps buffers; no longer needed
	RELEASE_COM(vsBuffer);
	RELEASE_COM(psBuffer);

	// Create texture sampler state desc
	D3D11_SAMPLER_DESC sd;
	ZeroMemory(&sd, sizeof(sd));

	sd.Filter = D3D11_FILTER_MIN_MAG_MIP_LINEAR;
	sd.AddressU = D3D11_TEXTURE_ADDRESS_WRAP;
	sd.AddressV = D3D11_TEXTURE_ADDRESS_WRAP;
	sd.AddressW = D3D11_TEXTURE_ADDRESS_WRAP;
	sd.MipLODBias = 0.0f;
	sd.MaxAnisotropy = 1;
	sd.ComparisonFunc = D3D11_COMPARISON_ALWAYS;
	sd.BorderColor[0] = 0.0f;
	sd.BorderColor[1] = 0.0f;
	sd.BorderColor[2] = 0.0f;
	sd.BorderColor[3] = 0.0f;
	sd.MinLOD = 0;
	sd.MaxLOD = D3D11_FLOAT32_MAX;

	HR(device->CreateSamplerState(&sd, &_pSamplerState));

	// Create matrix buffer
	D3D11_BUFFER_DESC mbd;
	ZeroMemory(&mbd, sizeof(mbd));
	mbd.Usage = D3D11_USAGE_DYNAMIC; // updated every frame
	mbd.ByteWidth = sizeof(MatrixBufferType);
	mbd.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
	mbd.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
	mbd.MiscFlags = 0;
	mbd.StructureByteStride = 0;

	HR(device->CreateBuffer(&mbd, NULL, &_pMatrixBuffer));

	// Create camera buffer
	D3D11_BUFFER_DESC cbd;
	ZeroMemory(&cbd, sizeof(cbd));

	cbd.Usage = D3D11_USAGE_DYNAMIC;
	cbd.ByteWidth = sizeof(CameraBufferType);
	cbd.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
	cbd.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
	cbd.MiscFlags = 0;
	cbd.StructureByteStride = 0;

	HR(device->CreateBuffer(&cbd, NULL, &_pCameraBuffer));

	// Create light buffer
	D3D11_BUFFER_DESC lbd;
	ZeroMemory(&lbd, sizeof(lbd));

	lbd.Usage = D3D11_USAGE_DYNAMIC;
	lbd.ByteWidth = sizeof(LightBufferType);
	lbd.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
	lbd.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
	lbd.MiscFlags = 0;
	lbd.StructureByteStride = 0;

	HR(device->CreateBuffer(&lbd, NULL, &_pLightBuffer));

	return true;
}

void Shader::Shutdown()
{
	RELEASE_COM(_pLightBuffer);
	RELEASE_COM(_pCameraBuffer);
	RELEASE_COM(_pMatrixBuffer);
	RELEASE_COM(_pSamplerState);
	RELEASE_COM(_pLayout);
	RELEASE_COM(_pPixelShader);
	RELEASE_COM(_pVertexShader);
}

bool Shader::Render(ID3D11DeviceContext * deviceContext,
					int indexCount,
					const XMMATRIX & worldMatrix,
					const XMMATRIX & viewMatrix,
					const XMMATRIX & projectionMatrix,
					ID3D11ShaderResourceView * texture,
					const XMFLOAT3 & lightDirection,
					const XMFLOAT4 & lightAmbientColor,
					const XMFLOAT4 & lightDiffuseColor,
					const XMFLOAT3 & cameraPosition,
					const XMFLOAT4 & specularColor,
					float specularPower)
{
	if(!SetShaderParameters(deviceContext,
						     worldMatrix,
							 viewMatrix,
							 projectionMatrix,
							 texture,
							 lightDirection,
							 lightAmbientColor,
							 lightDiffuseColor,
							 cameraPosition,
							 specularColor,
							 specularPower))
	{
		return false;
	}

	RenderShader(deviceContext, indexCount);

	return true;
}

void Shader::OutputShaderErrorMessage(ID3D10Blob * errorMsg,
									  HWND hwnd,
									  WCHAR *shaderFilename)
{
	char * compileErrors = (char *)errorMsg->GetBufferPointer();
	unsigned long bufferSize = errorMsg->GetBufferSize();

	std::ofstream fout;
	fout.open("shader-error.txt");

	for(unsigned i = 0; i < bufferSize; ++i)
	{
		fout << compileErrors[i];
	}

	fout.close();

	errorMsg->Release();

	MessageBox(hwnd,
			   L"Error compiling shader. Check shader-error.txt",
			   shaderFilename,
			   MB_OK);
}


bool Shader::SetShaderParameters(ID3D11DeviceContext *deviceContext,
								 const XMMATRIX & worldM,
							     const XMMATRIX & viewM,
								 const XMMATRIX & projectionM,
								 ID3D11ShaderResourceView * texture,
								 const XMFLOAT3 & lightDirection,
								 const XMFLOAT4 & lightAmbientColor,
								 const XMFLOAT4 & lightDiffuseColor,
								 const XMFLOAT3 & cameraPosition,
								 const XMFLOAT4 & specularColor,
								 float specularPower)
{
	// Transpose matrices to prepare them for shader: XMMATRIX is
	// row major while HLSL uniform parameters are column-major by default
	// Add #pragma pack_matrix(row_major) on top of HLSL files to change that
	XMMATRIX worldMatrix = XMMatrixTranspose(worldM);
	XMMATRIX viewMatrix = XMMatrixTranspose(viewM);
	XMMATRIX projectionMatrix = XMMatrixTranspose(projectionM);

	D3D11_MAPPED_SUBRESOURCE mappedResource;
	ZeroMemory(&mappedResource, sizeof(mappedResource));
	unsigned int bufferPos = 0; // constant buffer position in respective shader

	//  Lock constant buffer so it can be written to
	HR(deviceContext->Map(_pMatrixBuffer,
						  0,
						  D3D11_MAP_WRITE_DISCARD,
						  0,
						  &mappedResource));

	// Get pointer to data in constant buffer...
	MatrixBufferType * matrixBuf = (MatrixBufferType *) mappedResource.pData;
	// ... and fill it
	matrixBuf->world = worldMatrix;
	matrixBuf->view = viewMatrix;
	matrixBuf->projection = projectionMatrix;

	// Unlock
	deviceContext->Unmap(_pMatrixBuffer, 0);

	// Set constant buffer in vertex shader with updated values
	deviceContext->VSSetConstantBuffers(bufferPos, 1, &_pMatrixBuffer);

	// Lock camera buffer
	HR(deviceContext->Map(_pCameraBuffer, 0, D3D11_MAP_WRITE_DISCARD, 0, &mappedResource));
	CameraBufferType * cameraBuf = (CameraBufferType *)mappedResource.pData;

	// Copy camera position
	cameraBuf->cameraPosition = cameraPosition;

	// Unlock
	deviceContext->Unmap(_pCameraBuffer, 0);

	// Set camera constant buffer in vertex shader
	++bufferPos;
	deviceContext->VSSetConstantBuffers(bufferPos, 1, &_pCameraBuffer);

	HR(deviceContext->Map(_pLightBuffer,
						  0,
						  D3D11_MAP_WRITE_DISCARD,
						  0,
						  &mappedResource));

	LightBufferType * lightBuf = (LightBufferType *) mappedResource.pData;
	lightBuf->ambientColor = lightAmbientColor;
	lightBuf->diffuseColor = lightDiffuseColor;
	lightBuf->direction = lightDirection;
	lightBuf->specularColor = specularColor;
	lightBuf->specularPower = specularPower;

	deviceContext->Unmap(_pLightBuffer, 0);

	// Set light constant buffer position in pixel shader
	bufferPos = 0;
	deviceContext->PSSetConstantBuffers(bufferPos, 1, &_pLightBuffer);

	// Set shader texture resource in the pixel shader
	deviceContext->PSSetShaderResources(0, 1, &texture);

	return true;
}

void Shader::RenderShader(ID3D11DeviceContext * deviceContext, int indexCount)
{
	// Set vertex input layout
	deviceContext->IASetInputLayout(_pLayout);

	// Set vertex and pixel shaders
	deviceContext->VSSetShader(_pVertexShader, NULL, 0);
	deviceContext->PSSetShader(_pPixelShader, NULL, 0);

	// Set sampler state in pixel shader
	deviceContext->PSSetSamplers(0, 1, &_pSamplerState);

	deviceContext->DrawIndexed(indexCount, 0, 0);
}

}